#!/usr/bin/env python3
"""
高级数据优化工具
提供多种优化策略和自定义选项
"""

import pandas as pd
import numpy as np
import os
import json
import logging
from datetime import datetime
from utils import optimize_analysis_data, export_optimized_analysis

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class AdvancedDataOptimizer:
    """高级数据优化器"""
    
    def __init__(self):
        self.optimization_strategies = {
            'trading_focused': {
                'name': '交易导向',
                'description': '专注于交易信号和风险管理',
                'columns': [
                    '时间', '开盘价', '最高价', '最低价', '收盘价', '成交量',
                    'RSI相对强弱指数', 'MACD快慢线差值', 'MACD信号线',
                    '布林带上轨', '布林带下轨', '价格布林带位置',
                    'KDJ_K值', 'KDJ_D值', 'KDJ金叉', 'KDJ死叉',
                    '多头信号', '空头信号', '多头强度', '空头强度',
                    '多头止损位', '空头止损位', '建议仓位大小',
                    '风险等级', '交易建议', '综合评分'
                ]
            },
            'technical_analysis': {
                'name': '技术分析',
                'description': '专注于技术指标分析',
                'columns': [
                    '时间', '开盘价', '最高价', '最低价', '收盘价', '成交量',
                    '5日均线', '10日均线', '20日均线', '60日均线',
                    'RSI相对强弱指数', 'RSI超买信号', 'RSI超卖信号',
                    'MACD快慢线差值', 'MACD信号线', 'MACD柱状图', 'MACD金叉死叉',
                    '布林带上轨', '布林带中轨', '布林带下轨', '布林带宽度',
                    'KDJ_K值', 'KDJ_D值', 'KDJ_J值',
                    '真实波动幅度', 'ATR百分比', '平均趋向指数',
                    '商品通道指数', '威廉指标', '20日历史波动率'
                ]
            },
            'risk_management': {
                'name': '风险管理',
                'description': '专注于风险评估和仓位管理',
                'columns': [
                    '时间', '开盘价', '最高价', '最低价', '收盘价', '成交量',
                    '真实波动幅度', 'ATR百分比', '20日历史波动率',
                    '波动率分位数', '波动率突破',
                    '多头止损位', '空头止损位', '建议仓位大小', '最大仓位',
                    '风险评分', '风险等级', '趋势评分', '动量评分',
                    '恐慌贪婪指数', '市场情绪', '市场强度',
                    '综合评分', '交易建议'
                ]
            },
            'volume_analysis': {
                'name': '成交量分析',
                'description': '专注于成交量和资金流分析',
                'columns': [
                    '时间', '开盘价', '最高价', '最低价', '收盘价', '成交量',
                    '成交量加权均价', 'VWAP偏离度', '能量潮指标', 'OBV信号',
                    '资金流量指数', 'MFI超买', 'MFI超卖',
                    '5日历史波动率', '20日历史波动率',
                    '市场强度', '趋势强度'
                ]
            },
            'pattern_recognition': {
                'name': '形态识别',
                'description': '专注于K线形态和背离信号',
                'columns': [
                    '时间', '开盘价', '最高价', '最低价', '收盘价', '成交量',
                    'RSI相对强弱指数', 'MACD快慢线差值',
                    '顶背离', '底背离', '锤子线', '上吊线',
                    '突破多头', '突破空头', '布林带突破',
                    'KDJ金叉', 'KDJ死叉', 'MACD金叉死叉',
                    '交易建议', '趋势强度'
                ]
            }
        }
    
    def get_strategy_info(self):
        """获取所有优化策略信息"""
        return {key: {'name': value['name'], 'description': value['description']} 
                for key, value in self.optimization_strategies.items()}
    
    def optimize_with_strategy(self, df, strategy_name):
        """使用指定策略优化数据"""
        if strategy_name not in self.optimization_strategies:
            raise ValueError(f"未知策略: {strategy_name}")
        
        strategy = self.optimization_strategies[strategy_name]
        target_columns = strategy['columns']
        
        # 过滤存在的列
        available_columns = [col for col in target_columns if col in df.columns]
        missing_columns = [col for col in target_columns if col not in df.columns]
        
        if missing_columns:
            logger.warning(f"策略 {strategy['name']} 中以下列不存在: {missing_columns}")
        
        # 创建优化后的数据
        optimized_df = df[available_columns].copy()
        
        # 应用数据质量优化
        optimized_df = self._optimize_data_quality(optimized_df)
        
        logger.info(f"使用策略 '{strategy['name']}' 优化完成")
        logger.info(f"数据规模: {len(optimized_df)} 行 × {len(optimized_df.columns)} 列")
        
        return optimized_df
    
    def _optimize_data_quality(self, df):
        """优化数据质量"""
        # 处理数值精度
        numeric_columns = df.select_dtypes(include=[np.number]).columns
        for col in numeric_columns:
            if col not in ['时间', '成交量']:
                if any(keyword in col for keyword in ['价', '均线', 'MACD', '布林带', '止损', 'VWAP']):
                    df[col] = df[col].round(2)
                elif any(keyword in col for keyword in ['率', '百分比', 'RSI', 'KDJ', '波动', '评分']):
                    df[col] = df[col].round(4)
                else:
                    df[col] = df[col].round(3)
        
        # 优化布尔值列
        boolean_columns = [col for col in df.columns if any(keyword in col for keyword in 
                          ['信号', '金叉', '死叉', '超买', '超卖', '突破', '背离'])]
        for col in boolean_columns:
            if df[col].dtype == bool:
                df[col] = df[col].astype(int)
        
        # 优化分类列
        categorical_columns = ['风险等级', '交易建议', '市场情绪']
        for col in categorical_columns:
            if col in df.columns:
                df[col] = df[col].astype('category')
        
        return df
    
    def create_custom_optimization(self, df, custom_columns):
        """创建自定义优化"""
        available_columns = [col for col in custom_columns if col in df.columns]
        missing_columns = [col for col in custom_columns if col not in df.columns]
        
        if missing_columns:
            logger.warning(f"自定义优化中以下列不存在: {missing_columns}")
        
        optimized_df = df[available_columns].copy()
        optimized_df = self._optimize_data_quality(optimized_df)
        
        return optimized_df
    
    def compare_strategies(self, df):
        """比较不同策略的优化效果"""
        original_size = len(df) * len(df.columns)
        
        comparison = {
            'original': {
                'rows': len(df),
                'columns': len(df.columns),
                'data_points': original_size
            },
            'strategies': {}
        }
        
        for strategy_name, strategy in self.optimization_strategies.items():
            try:
                optimized_df = self.optimize_with_strategy(df, strategy_name)
                optimized_size = len(optimized_df) * len(optimized_df.columns)
                
                comparison['strategies'][strategy_name] = {
                    'name': strategy['name'],
                    'description': strategy['description'],
                    'rows': len(optimized_df),
                    'columns': len(optimized_df.columns),
                    'data_points': optimized_size,
                    'compression_ratio': round((1 - optimized_size / original_size) * 100, 1)
                }
            except Exception as e:
                logger.error(f"策略 {strategy_name} 优化失败: {e}")
        
        return comparison

def interactive_advanced_optimization():
    """交互式高级优化"""
    print("🚀 高级数据优化工具")
    print("=" * 60)
    
    optimizer = AdvancedDataOptimizer()
    
    # 选择文件
    analysis_files = []
    for filename in os.listdir('.'):
        if filename.endswith('_analysis.csv'):
            analysis_files.append(filename)
    
    if not analysis_files:
        print("❌ 未找到analysis.csv文件")
        return
    
    print(f"\n找到 {len(analysis_files)} 个分析文件:")
    for i, filename in enumerate(analysis_files, 1):
        print(f"{i}. {filename}")
    
    while True:
        try:
            choice = int(input(f"\n请选择文件 (1-{len(analysis_files)}): "))
            if 1 <= choice <= len(analysis_files):
                selected_file = analysis_files[choice - 1]
                break
            else:
                print("无效选择，请重新输入")
        except ValueError:
            print("请输入数字")
    
    # 加载数据
    try:
        df = pd.read_csv(selected_file)
        print(f"\n✅ 成功加载文件: {selected_file}")
        print(f"数据规模: {len(df)} 行 × {len(df.columns)} 列")
    except Exception as e:
        print(f"❌ 加载文件失败: {e}")
        return
    
    # 选择优化模式
    print(f"\n请选择优化模式:")
    print("1. 预设策略优化")
    print("2. 策略效果比较")
    print("3. 自定义列选择")
    
    while True:
        mode = input("\n请输入选择 (1-3): ").strip()
        if mode in ['1', '2', '3']:
            break
        else:
            print("无效选择，请重新输入")
    
    if mode == '1':
        # 预设策略优化
        strategies = optimizer.get_strategy_info()
        print(f"\n可用优化策略:")
        strategy_keys = list(strategies.keys())
        for i, (key, info) in enumerate(strategies.items(), 1):
            print(f"{i}. {info['name']} - {info['description']}")
        
        while True:
            try:
                choice = int(input(f"\n请选择策略 (1-{len(strategies)}): "))
                if 1 <= choice <= len(strategies):
                    selected_strategy = strategy_keys[choice - 1]
                    break
                else:
                    print("无效选择，请重新输入")
            except ValueError:
                print("请输入数字")
        
        # 执行优化
        try:
            optimized_df = optimizer.optimize_with_strategy(df, selected_strategy)
            
            # 保存结果
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            output_file = f"{selected_file.replace('_analysis.csv', '')}_{timestamp}_strategy_{selected_strategy}.csv"
            optimized_df.to_csv(output_file, index=False, encoding='utf-8-sig')
            
            print(f"\n✅ 优化完成!")
            print(f"输出文件: {output_file}")
            print(f"数据压缩: {(1 - len(optimized_df.columns) / len(df.columns)) * 100:.1f}%")
            
        except Exception as e:
            print(f"❌ 优化失败: {e}")
    
    elif mode == '2':
        # 策略效果比较
        print(f"\n🔍 正在比较所有策略效果...")
        comparison = optimizer.compare_strategies(df)
        
        print(f"\n📊 策略比较结果:")
        print(f"原始数据: {comparison['original']['rows']} 行 × {comparison['original']['columns']} 列")
        print("-" * 80)
        
        for strategy_name, info in comparison['strategies'].items():
            print(f"{info['name']:15} | {info['columns']:3d}列 | 压缩率: {info['compression_ratio']:5.1f}% | {info['description']}")
        
        # 保存比较报告
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        report_file = f"strategy_comparison_{timestamp}.json"
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(comparison, f, ensure_ascii=False, indent=2)
        
        print(f"\n📄 比较报告已保存: {report_file}")
    
    elif mode == '3':
        # 自定义列选择
        print(f"\n📋 当前数据包含以下列:")
        for i, col in enumerate(df.columns, 1):
            print(f"{i:3d}. {col}")
        
        print(f"\n请输入要保留的列编号 (用逗号分隔，如: 1,2,5-10,15):")
        selection = input("列编号: ").strip()
        
        try:
            # 解析选择
            selected_indices = []
            for part in selection.split(','):
                if '-' in part:
                    start, end = map(int, part.split('-'))
                    selected_indices.extend(range(start, end + 1))
                else:
                    selected_indices.append(int(part))
            
            # 转换为列名
            selected_columns = [df.columns[i-1] for i in selected_indices if 1 <= i <= len(df.columns)]
            
            if not selected_columns:
                print("❌ 没有选择有效的列")
                return
            
            # 执行自定义优化
            optimized_df = optimizer.create_custom_optimization(df, selected_columns)
            
            # 保存结果
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            output_file = f"{selected_file.replace('_analysis.csv', '')}_{timestamp}_custom.csv"
            optimized_df.to_csv(output_file, index=False, encoding='utf-8-sig')
            
            print(f"\n✅ 自定义优化完成!")
            print(f"输出文件: {output_file}")
            print(f"保留列数: {len(selected_columns)}")
            print(f"数据压缩: {(1 - len(selected_columns) / len(df.columns)) * 100:.1f}%")
            
        except Exception as e:
            print(f"❌ 自定义优化失败: {e}")

if __name__ == "__main__":
    interactive_advanced_optimization()
